from skimage import io
import ipywidgets
import numpy as np
from skimage import img_as_float
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.restoration import (denoise_tv_chambolle, denoise_bilateral, denoise_wavelet, denoise_nl_means)
import matplotlib.pyplot as plt


def fft2_denoise(im):
     #TODO 完成fft2去噪算法，实现输入为原图，输出为去噪后的图像
    h, w = im.shape
    fCoef = np.fft.fft2(im)
    amp_im_fft = np.log(np.abs(fCoef))
    keep = 0.1  # 每一个角落保留1/10的内容
    mask = np.ones((h, w))
    h_min = int(h * keep)
    h_max = int(h * (1 - keep))  # 需要掩盖住部分的纵坐标两端
    w_min = int(w * keep)
    w_max = int(w * (1 - keep))  # 需要掩盖住部分的横坐标两端
    mask[h_min:h_max, :] = 0
    mask[:, w_min:w_max] = 0
    fCoef_mask = fCoef * mask
    im_denoised = np.fft.ifft2(fCoef_mask).real
    return im_denoised


def optimize_denoise(img_list):
    # 以下虽然用了as_gray=True, 但是对于某些tif仍然无法转化为float
    # 所以仍然要使用img_as_floatpywidgets
    clean_img = img_as_float(io.imread(img_list[0], as_gray=True))  # 干净原图
    noisy_img = img_as_float(io.imread(img_list[1], as_gray=True))  # 有噪音图像
    # TODO 完成此部分代码实现题目要求功能

    # 1. 双边滤波去噪
    bilateral_img = denoise_bilateral(noisy_img, multichannel=False)
    # 2. 全变分去噪
    tv_img = denoise_tv_chambolle(noisy_img)
    # 3. 小波去噪
    wavelet_img = denoise_wavelet(noisy_img)
    # 4. Skimage的非局部均值去噪
    NLM_skimg_denoise_img = denoise_nl_means(noisy_img, patch_size=9, patch_distance=5)

    # 计算所有的去噪算法的PSNR（峰值信噪比）
    bilateral_psnr = psnr(clean_img, bilateral_img)
    tv_psnr = psnr(clean_img, tv_img)
    wavelet_psnr = psnr(clean_img, wavelet_img)
    NLM_psnr = psnr(clean_img, NLM_skimg_denoise_img)
    fft2_psnr = psnr(clean_img, fft2_denoise(noisy_img))

    print("bilateral PSNR:", bilateral_psnr)
    print("tv PSNR:", tv_psnr)
    print("wavelet PSNR:", wavelet_psnr)
    print("NLM PSNR:", NLM_psnr)
    print("fft2 PSNR:", fft2_psnr)

    psnr_list = np.array([bilateral_psnr, tv_psnr, wavelet_psnr, NLM_psnr, fft2_psnr])
    i = np.argmax(psnr_list)
    img_list = np.array([bilateral_img, tv_img, wavelet_img, NLM_skimg_denoise_img, fft2_denoise(noisy_img)])
    algorithm_list = np.array(['bilateral', 'tv', 'wavelet', 'NLM', 'fft2'])

    fg, ax = plt.subplots(1, 3, figsize=(10, 4))
    ax[0].imshow(clean_img, cmap="gray")
    ax[0].set_title("clean_img")
    ax[1].imshow(noisy_img, cmap="gray")
    ax[1].set_title("noisy_img")
    ax[2].imshow(img_list[i], cmap="gray")
    ax[2].set_title("{0}:{1}".format(algorithm_list[i], psnr_list[i]))
    plt.tight_layout()
    plt.show()


x = ipywidgets.Dropdown(
    options=[('MRI', ("img/MRI_clean.tif", "img/MRI_noisy.tif")),
             ('cat', ("img/cat.png", "img/cat_noisy.jpg")),
             ('balloon', ("img/balloon.jpg", "img/balloon_noisy.jpg")),
             ('einstein', ("img/einstein.jpg", "img/einstein_noisy.jpg")),
             ]
)
ipywidgets.interact(optimize_denoise, img_list=x)

